import numpy as np

import sys
import os 


from .base import  get_jacobian, get_measurement



class heading_filter_loosely_couple:
    def __init__(self, d):
        self.theta = np.array([-0.0, 0.0, 0.0])
        self.P_prev = np.eye(3)
        self.Q =  np.eye(3)
        self.R =  np.eye(3)
        self.d = d
    
    def predict(self):
        F = np.eye(3)
        self.theta = F @ self.theta
        self.P_pred = F @ self.P_prev @ F.T + self.Q
    
    def update(self, delta_p):
        s_1, s_2, s_3 = np.sin(self.theta[0]), np.sin(self.theta[1]), np.sin(self.theta[2])
        c_1, c_2, c_3 = np.cos(self.theta[0]), np.cos(self.theta[1]), np.cos(self.theta[2])

        # Define the matrix
        H = np.array([[s_1*s_3+c_1*s_2*c_3, s_1*c_2*c_3, -c_1*c_3-s_1*s_2*s_3],
                        [-s_1*c_3+c_1*s_2*s_3, s_1*c_2*s_3, -c_1*s_3+s_1*s_2*c_3],
                        [c_1*c_2, -s_1*s_2, 0]])*self.d
        K = self.P_pred @ H.T @ np.linalg.inv(H @ self.P_pred @ H.T + self.R)
    
        self.theta = self.theta + K @ (delta_p - H @ self.theta)
        self.P_est = (np.eye(3) - K @ H) @ self.P_pred


class heading_filter_tightly_couple:
    def __init__(self, anchors, tags, state_init):
        self.state_size = 4
        self.state = state_init
        self.P = np.eye(self.state_size)
        self.P[0:3, 0:3] *= 0.1
        self.P[3, 3] *= 0.01
        self.Q =  self.P
        self.roll = 0
        self.pitch = 0
        self.anchors = anchors
        self.tags = tags
    def func(self):
        return get_measurement(self.state[0:3], np.array([self.roll, self.pitch, self.state[3]]), self.anchors, self.tags)
    def predict(self):
        F = np.eye(self.state_size)
        self.state = F @ self.state
        self.P = F @ self.P @ F.T + self.Q
    
    def update(self, z):
        # Define the matrix
        H = get_jacobian(self.anchors, self.tags, self.state[0:3], np.array([self.roll, self.pitch, self.state[3]]))
        R = np.eye(z.shape[0]) * 0.3
        K = self.P @ H.T @ np.linalg.inv(H @ self.P @ H.T + R)
    
        self.state = self.state + K @ (z - self.func())
        self.P = (np.eye(self.state_size) - K @ H) @ self.P